from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import IntegerType

if __name__ == '__main__':
    spark = SparkSession.builder. \
        appName("test"). \
        master("local[*]"). \
        config("spark.sql.shuffle.partitions", 2). \
        getOrCreate()

    sc = spark.sparkContext

    rdd = sc.parallelize([1, 2, 3, 4, 5, 6, 7]).map(lambda x: [x])
    df = rdd.toDF(["num"])


    # TODO 1：方式1 sparksession.udf.register()，DSL和SQL风格均可以使用
    def num_ride_10(num):
        return num * 10


    udf2 = spark.udf.register("udf1", num_ride_10, IntegerType())

    # SQL风格使用
    df.selectExpr("udf1(num)").show()

    # DSL风格
    df.select(udf2(df['num'])).show()

    # TODO 2: 方式2注册，仅用于DSL风格
    udf3 = F.udf(num_ride_10, IntegerType())
    df.select(udf3(df['num'])).show()
